import numpy as np
from sgfmill import sgf
from sgfmill import boards
import os
from typing import List, Tuple, Optional, Dict, Any

class SGFToNumpyConverter:
    """
    将SGF文件中的每一步棋转换为可训练的numpy数组
    """
    
    def __init__(self, board_size: int = 19):
        """
        初始化转换器
        
        Args:
            board_size: 棋盘大小，默认为19
        """
        self.board_size = board_size
        self.num_features = 17  # 特征平面数量
        
    def load_sgf_file(self, file_path: str) -> Optional[sgf.Sgf_game]:
        """
        加载SGF文件
        
        Args:
            file_path: SGF文件路径
            
        Returns:
            sgf.Sgf_game对象或None（如果加载失败）
        """
        try:
            with open(file_path, 'rb') as f:
                game_data = f.read()
            return sgf.Sgf_game.from_bytes(game_data)
        except Exception as e:
            print(f"加载SGF文件失败 {file_path}: {e}")
            return None
    
    def get_all_moves(self, game: sgf.Sgf_game) -> List[Tuple[Optional[str], Optional[Tuple[int, int]]]]:
        """
        获取游戏中的所有移动
        
        Args:
            game: sgf.Sgf_game对象
            
        Returns:
            移动列表，每个元素为(颜色, 坐标)
        """
        moves = []
        node = game.get_root()
        
        while node:
            try:
                color, move = node.get_move()
                moves.append((color, move))
                
                # 获取下一个节点
                if node:
                    node = node[0]
                else:
                    break
            except (IndexError, ValueError, AttributeError):
                break
        
        return moves
    
    def create_empty_board(self) -> np.ndarray:
        """创建空棋盘状态"""
        return np.zeros((self.board_size, self.board_size), dtype=np.int8)
    
    def apply_move_to_board(self, board: np.ndarray, color: str, move: Optional[Tuple[int, int]]) -> np.ndarray:
        """
        在棋盘上应用移动
        
        Args:
            board: 当前棋盘状态
            color: 棋子颜色 'b' 或 'w'
            move: 移动坐标，None表示pass
            
        Returns:
            更新后的棋盘状态
        """
        new_board = board.copy()
        
        if move is not None:
            row, col = move
            if 0 <= row < self.board_size and 0 <= col < self.board_size:
                # 将颜色转换为数字：黑棋为1，白棋为-1
                stone_value = 1 if color == 'b' else -1
                new_board[row, col] = stone_value
        
        return new_board
    
    def simulate_board_state(self, moves: List[Tuple[Optional[str], Optional[Tuple[int, int]]]]) -> List[np.ndarray]:
        """
        模拟整个游戏的棋盘状态
        
        Args:
            moves: 所有移动列表
            
        Returns:
            每一步对应的棋盘状态列表
        """
        board_states = []
        current_board = self.create_empty_board()
        
        for color, move in moves:
            if color is not None:
                board_states.append(current_board.copy())
                current_board = self.apply_move_to_board(current_board, color, move)
        
        return board_states
    
    def board_to_features(self, board: np.ndarray, current_player: str) -> np.ndarray:
        """
        将棋盘状态转换为特征平面
        
        Args:
            board: 棋盘状态数组
            current_player: 当前玩家 'b' 或 'w'
            
        Returns:
            形状为 (num_features, board_size, board_size) 的特征张量
        """
        features = []
        
        # 1. 当前玩家的棋子
        current_player_stones = np.zeros((self.board_size, self.board_size), dtype=np.float32)
        # 2. 对手的棋子
        opponent_stones = np.zeros((self.board_size, self.board_size), dtype=np.float32)
        
        current_value = 1 if current_player == 'b' else -1
        opponent_value = -1 if current_player == 'b' else 1
        
        current_player_stones[board == current_value] = 1.0
        opponent_stones[board == opponent_value] = 1.0
        
        features.extend([current_player_stones, opponent_stones])
        
        # 3. 空点
        empty_points = ((board == 0).astype(np.float32))
        features.append(empty_points)
        
        # 4-7. 最近4步的移动历史（简化实现）
        for i in range(4):
            history_plane = np.zeros((self.board_size, self.board_size), dtype=np.float32)
            features.append(history_plane)
        
        # 8. 当前玩家颜色平面
        player_color_plane = np.ones((self.board_size, self.board_size), dtype=np.float32)
        if current_player == 'w':
            player_color_plane *= 0.0
        features.append(player_color_plane)
        
        # 9-17. 常数平面（可根据需要扩展为更复杂的特征）
        for i in range(8):
            constant_plane = np.zeros((self.board_size, self.board_size), dtype=np.float32)
            features.append(constant_plane)
        
        # 添加一个全1的常数平面
        ones_plane = np.ones((self.board_size, self.board_size), dtype=np.float32)
        features.append(ones_plane)
        
        return np.stack(features, axis=0)
    
    def move_to_label(self, move: Optional[Tuple[int, int]]) -> int:
        """
        将移动转换为标签索引
        
        Args:
            move: 移动坐标，None表示pass
            
        Returns:
            标签索引（0到board_size*board_size）
        """
        if move is None:  # Pass
            return self.board_size * self.board_size
        
        row, col = move
        return row * self.board_size + col

    def normalize_labels(self, labels: np.ndarray) -> np.ndarray:
        """
        对标签进行归一化到[0, 1]范围
        
        Args:
            labels: 原始标签数组
            
        Returns:
            归一化后的标签数组
        """
        max_label = self.board_size * self.board_size
        return labels.astype(np.float32) / max_label

    def labels_to_one_hot(self, labels: np.ndarray) -> np.ndarray:
        """
        将标签转换为one-hot编码
        
        Args:
            labels: 原始标签数组
            
        Returns:
            one-hot编码的标签数组
        """
        num_classes = self.board_size * self.board_size + 1  # 包括pass
        return np.eye(num_classes, dtype=np.float32)[labels]
    
    def process_single_game(self, game: sgf.Sgf_game, label_format: str = 'original') -> Tuple[np.ndarray, np.ndarray]:
        """
        处理单个SGF游戏，转换为训练数据
        
        Args:
            game: sgf.Sgf_game对象
            label_format: 标签格式 'original', 'normalized', 或 'one_hot'
            
        Returns:
            (features, labels) 特征数组和标签数组
        """
        # 检查棋盘大小
        if game.get_size() != self.board_size:
            print(f"棋盘大小不匹配: 期望{self.board_size}, 实际{game.get_size()}")
            return np.array([]), np.array([])
        
        # 获取所有移动
        moves = self.get_all_moves(game)
        if not moves:
            return np.array([]), np.array([])
        
        # 过滤掉颜色为None的移动
        valid_moves = [(color, move) for color, move in moves if color is not None]
        if not valid_moves:
            return np.array([]), np.array([])
        
        # 模拟棋盘状态
        board_states = self.simulate_board_state(valid_moves)
        
        all_features = []
        all_labels = []
        
        for i, ((color, move), board_state) in enumerate(zip(valid_moves, board_states)):
            # 提取当前状态的特征
            features = self.board_to_features(board_state, color)
            all_features.append(features)
            
            # 创建标签
            label = self.move_to_label(move)
            all_labels.append(label)
        
        if all_features:
            features_array = np.array(all_features)
            labels_array = np.array(all_labels)
            
            # 根据选择的格式处理标签
            if label_format == 'minmax':
                labels_array = self.normalize_labels(labels_array)
            elif label_format == 'one_hot':
                labels_array = self.labels_to_one_hot(labels_array)
            # 'original' 保持不变
            
            #print(f"从单个游戏生成 {len(features_array)} 个样本")
            #print(f"标签格式: {label_format}")
            #print(f"特征形状: {features_array.shape}")
            #print(f"标签形状: {labels_array.shape}")
            
            return features_array, labels_array
        else:
            return np.array([]), np.array([])
    
    def save_training_data(self, features: np.ndarray, labels: np.ndarray, output_path: str):
        """
        保存训练数据为NPZ文件
        
        Args:
            features: 特征数组
            labels: 标签数组
            output_path: 输出文件路径
        """
        if len(features) == 0 or len(labels) == 0:
            print("没有数据可保存")
            return
        
        np.savez_compressed(
            output_path,
            features=features,
            labels=labels,
            board_size=self.board_size,
            num_features=self.num_features
        )
        print(f"训练数据已保存到: {output_path}")
        if os.path.exists(output_path):
            print(f"文件大小: {os.path.getsize(output_path) / 1024 / 1024:.2f} MB")
    
    def load_training_data(self, file_path: str) -> Tuple[np.ndarray, np.ndarray]:
        """
        加载训练数据
        
        Args:
            file_path: NPZ文件路径
            
        Returns:
            (features, labels) 特征和标签数组
        """
        try:
            data = np.load(file_path)
            features = data['features']
            labels = data['labels']
            #print(f"从 {file_path} 加载数据成功")
            #print(f"特征形状: {features.shape}")
            #print(f"标签形状: {labels.shape}")
            return features, labels
        except Exception as e:
            print(f"加载训练数据失败: {e}")
            return np.array([]), np.array([])

    def get_data_statistics(self, features: np.ndarray, labels: np.ndarray) -> Dict[str, Any]:
        """
        获取数据统计信息
        
        Args:
            features: 特征数组
            labels: 标签数组
            
        Returns:
            统计信息字典
        """
        if len(features) == 0:
            return {}
        
        stats = {
            'total_samples': len(features),
            'feature_shape': features[0].shape,
            'label_range': (labels.min(), labels.max()),
            'pass_moves': np.sum(labels == self.board_size * self.board_size),
            'unique_labels': len(np.unique(labels)),
            'data_type': features.dtype
        }
        
        # 如果是归一化标签，添加额外信息
        if labels.dtype == np.float32 and labels.max() <= 1.0:
            stats['label_type'] = 'minmax'
        elif len(labels.shape) > 1 and labels.shape[1] == self.board_size * self.board_size + 1:
            stats['label_type'] = 'one_hot'
        else:
            stats['label_type'] = 'original'
        
        return stats

# 简化使用函数
def convert_sgf_to_training_data(sgf_file: str, output_file: str, board_size: int = 19, label_format: str = 'original'):
    """
    简化函数：直接将SGF文件转换为训练数据文件
    
    Args:
        sgf_file: SGF文件路径
        output_file: 输出文件路径
        board_size: 棋盘大小
        label_format: 标签格式
    """
    converter = SGFToNumpyConverter(board_size=board_size)
    
    print("开始处理SGF文件...")
    game = converter.load_sgf_file(sgf_file)
    
    if game:
        features, labels = converter.process_single_game(game, label_format=label_format)
        
        if len(features) > 0:
            converter.save_training_data(features, labels, output_file)
            
            # 显示统计信息
            stats = converter.get_data_statistics(features, labels)
            print("\n数据统计:")
            for key, value in stats.items():
                print(f"  {key}: {value}")
        else:
            print("没有生成训练数据")
    else:
        print("无法加载SGF文件")

if __name__ == "__main__":
    # 使用示例
    converter = SGFToNumpyConverter(board_size=19)
    
    # 处理单个SGF文件，使用归一化标签
    game = converter.load_sgf_file("example.sgf")
    if game:
        # 使用原始标签
        features, labels = converter.process_single_game(game, label_format='original')
        
        # 使用归一化标签
        features_norm, labels_norm = converter.process_single_game(game, label_format='normalized')
        
        # 使用one-hot编码
        features_onehot, labels_onehot = converter.process_single_game(game, label_format='one_hot')
        
        if len(features) > 0:
            converter.save_training_data(features, labels, "single_game_data.npz")
    
    # 或者使用简化函数
    # convert_sgf_to_training_data("example.sgf", "training_data.npz", label_format='normalized')